面向AI图形的轻量深度学习推理引擎 ShaderNN

背景：

近年来，随着移动端算力的不断提升和深度学习研究的快速发展，特别是小网络模型不断成熟以及数据安全的要求越来越高，越来越多的原本在云端执行的推理转移到移动端上来实现。移动平台的深度学习推理涉及硬件平台、驱动、编译优化、模型压缩、算子算法优化以及部署，适合系统业务开发的高效推理框架成为业界迫切需求和开发重点。

很多企业都在研发面向移动端的开源深度学习框架。当前国内外主流开源移动端推理框架：小米MACE、腾讯优图NCNN/TNN、阿里MNN、谷歌TensorFlow Lite等。现有的推理平台大多依赖于不同的第三方库和硬件平台驱动，针对不同手机硬件和系统适配、验证、优化、部署、维护、升级繁琐。另一方面，受移动端软硬件限制，如有限的计算能力、严格的低功耗要求、有限的I/O带宽、不同的内存架构和访问方式等，对于移动端的主流图形图像类后处理的AI应用，例如视频实时超分、AI渲染、光追降噪、游戏AI后处理等，它们对实时性、性能和功耗要求比较极致，需要和应用在图形图像管线深度耦合，如需要直接基于GPU纹理进行推理，现有的推理框架很难达到要求。从硬件实现上，现有的推理框架通常使用单一硬件如CPU，GPU，NPU或者DSP进行推理，对于复杂场景和模型很难满足整体推理性能要求（如现有的硬件无法兼容、新的算子或模型不能支持，现有硬件上算子实现效能低等问题），需要一款轻量、可定制、可无缝嵌入应用并支持异构计算的框架来弥补以上不足。

基于移动端图形图像后处理高效AI推理需求，为降低业务融合成本，提高效能，我们开发了基于GPU Shader高效推理引擎ShaderNN。它直接基于GPU纹理进行高效推理以节省I/O时间， 不依赖第三方库，跨不同硬件平台，支持主流深度学习训练框架，可定制，方便优化、集成、部署和升级。

一、ShaderNN是什么？

图1:  ShaderNN在图形图像类AI应用

ShaderNN的全称是Shader Neural Network，顾名思义，它是基于图形Shader（着色器）的推理框架。如图1所示，在移动设备上图形和图像后处理的AI应用中，主要有基于OpenGL和基于Vulkan两种渲染管线。现有的大多数图形图像应用使用OpenGL渲染管线，ShaderNN提供基于OpenGL片段着色器（Fragment Shader）和计算着色器（Compute Shader）的深度学习模型推理实现，做到和OpenGL渲染管线的无缝衔接，支持主流的卷积神经网络模型和算子，这部分功能在第一阶段已开发完成并开源。越来越多新的应用、引擎和框架基于Vulkan开发，第二阶段将开发基于Vulkan渲染管线计算着色器（Compute Shader）的推理功能，很快推出并开源。

图2:  ShaderNN 工作流程图

如图2所示，ShaderNN 整个工作流程可以分为模型转换和层融合优化、网络模型和权重加载解析、计算图生成、计算图算子的执行、推理预测返回相关结果。ShaderNN支持通过TensorFlow、PyTorch训练的模型和转换成ONNX格式的模型,将它们转换成ShaderNN JSON格式。在转换过程中，模型转换器解耦模型结构和权重，解析算子并做层间融合优化。在加载模结构和权重以后，推理引擎通过拓扑排序生成计算图。并交由Shader来执行具体的算子。推理引擎分别在编译阶段和运行阶段进行优化。编译的优化包括Shader的编译和缓存优化，等效算子替换与融合等。在运行阶段，ShaderNN通过卷积优化、纹理重用、CPU和GPU 内存重用、数据结构布局优化、缓存及向量化优化以提升效能。

二、ShaderNN的优势

(1) 面向实时图形图像后处理应用：ShaderNN 直接操作于图形和图像应用程序的纹理数据，可以和图形和图像渲染管道紧密耦合, 这对于产品落地至关重要。对于实时要求高的AI应用， 性能的开销主要来自两个方面，一个是推理本身的开销，另外一个是推理框架和现有管线的I/O交互开销。后者一开始容易被忽略，却成为移动端AI产品应用落地的瓶颈。图3对比了ShaderNN和其它推理引擎与现有管线集成的不同，其它推理引擎通常提供了称为 Tensor 或 Mat 的通用数据结构作为输入和输出的格式，无法使用广泛用于图形和图像应用程序的纹理作为输入输出，需要进行格式转换或者数据拷贝。对于图形图像类的大输入的实时应用，这部分的开销是相当大的。而ShaderNN以GPU中的原生纹理作为输入输出，省去了数据格式转换和移动的开销，大大提升了应用的整体性能。

图3:  具有深度学习模型的图形和图像处理管道

(2) 高性能：ShaderNN是业内第一个基于集成片段着色器（Fragment Shader）和计算着色器（Compute Shader）混合方式实现的静态计算图推理引擎，这样的实现方式可以在运行前对计算图的结构进行优化，在典型使用场景下兼顾计算效率和灵活性，获得更快的推理速度。

(3) 轻量性、可移植性和可扩展性：ShaderNN推理核心库文件大小小于2M，无第三方依赖库，可方便部署到移动设备上。同时支持可定制算子，方便支持新的模型、算子及其优化。

(4) 通用性：支持TensorFlow，PyTorch和ONNX主流训练框架和模型格式。支持常用基于CNN的分类，目标识别、图像分割和图像增强网络结构，并提供ModelZoo和Android App方便开发者参考。

三、 性能和功耗

我们挑选了Spatial Denoise， ESPCN， Resnet18，YOLO V3 Tiny 4种CNN常用模型，分别在两款联发科和两款高通平台（见图4）上，与TensorFlow Lite OpenGL Benchmark 做了性能和功耗的对比：

图4:  测试移动设备

1、ShaderNN 和TensorFlow Lite 性能对比

图5:  性能对比

从图5性能对比上看，ShaderNN在推理时间上都有显著的提升。在Spatial Denoise和ESPCN方面优于 TensorFlow Lite 75%-90%；Resnet18 和 YOLO V3 Tiny模型在部分处理器芯片上优于TensorFlow Lite 50% 。

2、ShaderNN 和 TensorFlow Lite 功耗对比

图6:  功耗对比

从图6功耗对比上看，ShaderNN单次推理功耗明显少于TensorFlow Lite，平均功耗也小于TensorFlow Lite。

与TensorFlow Lite对比，ShaderNN 在选定的目标设备上与Spatial Denoise、ESPCN、Resnet18 和 YOLO V3 Tiny对比时，可节省高达 80%、70%、55% 和 51% 的功耗。

四、 ShaderNN的典型应用场景

从上面性能和功耗对比可以看出，ShaderNN在图形和图像类应用中有明显的优势。典型的应用场景如图7所示：光线追踪降噪、深度学习超采样、高动态范围成像（HDR）、超分辨率和风格迁移 [1]等。

图7:  ShaderNN典型场景

图8：展示4风格迁移的Android应用

五、规划和展望

作为一款基于GPU Shader的轻量、高效开源移动端推理引擎， ShaderNN通过和图形渲染管线的深度耦合，在实时图形图像处理AI应用中具有明显优势。随着移动端GPU的性能提升，更多的AI应用在移动端部署成为可能。我们希望基于GPU Shader的推理引擎能给业界带来新的动力, 优化落地更多产品，推动更多以前只能在桌面端实现的应用，如DLSS（Deep Learning Super Sampling， 深度学习超级采样）， Ray Tracing Denoising （光追降噪）等在移动端的实现，探索移动端GPU和NPU/DSP如何形成有效互补。

在基于OpenGL的ShaderNN推理框架作为开源项目发布后，我们会持续添加基于Vulkan的推理功能，覆盖更多的应用场景。为开源社区做出贡献的同时，也希望能吸引更多的开发者参与进来。我们欢迎其他移动开发人员来使用和共建这个项目，在下面列出的GitHub开源地址里和我们交流，一起支持更多的模型和算子，持续优化性能。我们将以合作共赢为使命，推动开源项目的发展，致力于共同打造一个基于OpenGL、Vulkan的全图形栈的移动推理引擎，部署到更多移动GPU平台产品中。

六、总结

本文介绍了一款开源的轻量级卷积神经网络的推理引擎ShaderNN。它基于GPU shader实现，面向对实时要求高的图形图像处理AI应用，通过和图形渲染OpenGL和Vulkan管线的深度耦合，提供高性能低功耗的推理能力，方便集成部署到移动平台。

GitHub开源地址 （Apache2.0开源许可）

https://github.com/inferenceengine/shadernn

参考文献：

1、Justin Johnson, Alexandre Alahi, and Li Fei-Fei. 2016. Perceptual losses for real-time style transfer and super-resolution. In European conference on computer vision. Springer, 694–711

长按关注内核工匠微信

Linux 内核黑科技 | 技术文章 | 精选教程